﻿using System;
using System.Collections.Generic;
using System.IO;
using System.Net.Http;
using System.Text;
using System.Text.Json;
using System.Threading.Tasks;

public class ImageParameters
{
    public string size { get; set; } = "1024*1024";
    public int n { get; set; } = 1;
    public bool prompt_extend { get; set; } = true;
    public bool watermark { get; set; } = false;
}

class Program
{
    private const string DEFAULT_API_KEY = "sk-817d27e9c883406a87d98eb23fc4b1a9";
    private static readonly string baseUrl = "https://dashscope.aliyuncs.com";

    static async Task Main(string[] args)
    {
        Console.WriteLine("Wanx2.1-t2i-turbo 图像生成器");

        string apiKey = Environment.GetEnvironmentVariable("DASHSCOPE_API_KEY") ?? DEFAULT_API_KEY;
        if (apiKey == DEFAULT_API_KEY)
        {
            Console.WriteLine("警告: 使用默认API Key，建议设置环境变量: export DASHSCOPE_API_KEY=\"your_api_key\"");
        }

        Console.WriteLine("请输入提示词 (例如: 雪地，白色小教堂，极光，冬日场景，柔和的光线):");
        string prompt = Console.ReadLine();

        Console.WriteLine("请输入负面提示词 (可选，例如: 人物，按Enter跳过):");
        string negativePrompt = Console.ReadLine();

        Console.WriteLine("请输入照片大小 (可选，默认1024*1024，按Enter跳过):");
        string size = Console.ReadLine();
        if (string.IsNullOrWhiteSpace(size)) size = "1024*1024";

        Console.WriteLine("请输入照片数量 (1-4，默认1，按Enter跳过):");
        string nInput = Console.ReadLine();
        int n = string.IsNullOrWhiteSpace(nInput) ? 1 : Math.Clamp(int.Parse(nInput), 1, 4);

        ImageParameters parameters = new ImageParameters
        {
            size = size,
            n = n,
            prompt_extend = true,
            watermark = false
        };

        Console.WriteLine("正在生成图像...");

        try
        {
            string taskId = await SubmitImageGenerationTask(apiKey, prompt, negativePrompt, parameters);
            Console.WriteLine($"任务已提交，任务ID: {taskId}");
            Console.WriteLine("正在等待任务完成...");

            List<string> imageUrls = await WaitForTaskCompletion(apiKey, taskId);
            Console.WriteLine($"\n成功生成 {imageUrls.Count} 张图片!");

            // 创建下载目录
            string downloadDir = $"生成结果_{DateTime.Now:yyyyMMdd_HHmmss}";
            Directory.CreateDirectory(downloadDir);

            // 下载所有图片
            for (int i = 0; i < imageUrls.Count; i++)
            {
                string filename = Path.Combine(downloadDir, $"image_{i + 1}.png");
                await DownloadImage(imageUrls[i], filename);
                Console.WriteLine($"图片 {i + 1} 已保存到: {filename}");
            }

            Console.WriteLine($"\n所有图片已保存到目录: {Path.GetFullPath(downloadDir)}");
        }
        catch (Exception ex)
        {
            Console.WriteLine($"发生错误: {ex.Message}");
        }

        Console.WriteLine("按任意键退出...");
        Console.ReadKey();
    }

    static async Task<string> SubmitImageGenerationTask(string apiKey, string prompt, string negativePrompt, ImageParameters p)
    {
        string apiUrl = $"{baseUrl}/api/v1/services/aigc/text2image/image-synthesis";

        var requestData = new
        {
            model = "wanx2.1-t2i-turbo",
            input = new
            {
                prompt = prompt,
                negative_prompt = string.IsNullOrEmpty(negativePrompt) ? null : negativePrompt
            },
            parameters = p
        };

        using var client = new HttpClient();
        client.DefaultRequestHeaders.Add("Authorization", $"Bearer {apiKey}");
        client.DefaultRequestHeaders.Add("X-DashScope-Async", "enable");
        client.DefaultRequestHeaders.Add("Accept", "application/json");

        var content = new StringContent(
            JsonSerializer.Serialize(requestData, new JsonSerializerOptions { PropertyNamingPolicy = JsonNamingPolicy.CamelCase }),
            Encoding.UTF8,
            "application/json");

        var response = await client.PostAsync(apiUrl, content);
        response.EnsureSuccessStatusCode();

        using JsonDocument doc = JsonDocument.Parse(await response.Content.ReadAsStringAsync());
        if (doc.RootElement.TryGetProperty("output", out var output) &&
            output.TryGetProperty("task_id", out var taskId))
        {
            return taskId.GetString();
        }
        throw new Exception("无法从响应中获取任务ID");
    }

    static async Task<List<string>> WaitForTaskCompletion(string apiKey, string taskId, int maxRetries = 20, int delaySeconds = 5)
    {
        string apiUrl = $"{baseUrl}/api/v1/tasks/{taskId}";

        using var client = new HttpClient();
        client.DefaultRequestHeaders.Add("Authorization", $"Bearer {apiKey}");
        client.DefaultRequestHeaders.Add("Accept", "application/json");

        for (int i = 0; i < maxRetries; i++)
        {
            Console.Write(".");
            var response = await client.GetAsync(apiUrl);
            response.EnsureSuccessStatusCode();

            using JsonDocument doc = JsonDocument.Parse(await response.Content.ReadAsStringAsync());
            JsonElement root = doc.RootElement;

            if (root.TryGetProperty("output", out var output))
            {
                string status = output.GetProperty("task_status").GetString();

                if (status == "SUCCEEDED")
                {
                    var imageUrls = new List<string>();
                    if (output.TryGetProperty("results", out var results))
                    {
                        foreach (var result in results.EnumerateArray())
                        {
                            if (result.TryGetProperty("url", out var url))
                            {
                                imageUrls.Add(url.GetString());
                            }
                        }
                    }
                    return imageUrls.Count > 0 ? imageUrls : throw new Exception("任务成功但未获取到任何图片URL");
                }
                else if (status == "FAILED")
                {
                    string code = output.GetProperty("code").GetString();
                    string message = output.GetProperty("message").GetString();
                    throw new Exception($"任务处理失败\ncode: {code}\nmessage: {message}");
                }
            }
            await Task.Delay(delaySeconds * 1000);
        }
        throw new Exception("任务处理超时");
    }

    static async Task DownloadImage(string imageUrl, string savePath)
    {
        using var client = new HttpClient();
        using var response = await client.GetAsync(imageUrl);
        response.EnsureSuccessStatusCode();

        using var fileStream = File.Create(savePath);
        await response.Content.CopyToAsync(fileStream);
    }
}